#include "fermiqcd.h"

void generate_stochastic(fermi_field& phi,
			 gauge_field& U,
			 coefficients& quark,
			 float absolute_precision,
			 float relative_precision) {
  int nc=phi.nc;
  fermi_field dummy(phi.lattice(),nc);
  set_random(dummy);
  mul_invQ(phi,dummy,U,quark,absolute_precision,relative_precision);
}



void build_shifted(fermi_field &sink_n_shifted, gauge_field &U, fermi_field &sink_n, int ipath) {
  int mu,nu,i,j,a;
  mdp_complex A1[3][3], A2[3][3], A3[3][3], A4[3][3];
  mdp_site x(sink_n_shifted.lattice());
  mdp_site y1(sink_n_shifted.lattice());
  mdp_site y2(sink_n_shifted.lattice());
  mdp_site y3(sink_n_shifted.lattice());
  mdp_site y4(sink_n_shifted.lattice());
     
  if(ipath==0) {

    sink_n_shifted=sink_n;

  } else if(ipath<5) {

    mu=ipath-1;
    forallsites(x)
      for(a=0; a<4; a++)
	for(i=0; i<3; i++)
	  for(j=0; j<3; j++)
	    sink_n_shifted(x,a,i)=0.5*(U(x,+1,mu,i,j)*sink_n(x+mu,a,j)-
				       U(x,-1,mu,i,j)*sink_n(x-mu,a,j));	
  } else {

    mu=(ipath-5)/4;
    nu=(ipath-5)%4;
    forallsites(x) {
      y1=(x+mu)+nu;
      y2=(x+mu)-nu;
      y3=(x-mu)+nu;
      y4=(x-mu)-nu;
      for(i=0; i<3; i++)
	for(j=0; j<3; j++) {
	  A1[i][j]=U(x,+1,mu,i,0)*U(x+mu,+1,nu,0,j)+U(x,+1,mu,i,1)*U(x+mu,+1,nu,1,j)+U(x,+1,mu,i,2)*U(x+mu,+1,nu,2,j);
	  A2[i][j]=U(x,+1,mu,i,0)*U(x+mu,-1,nu,0,j)+U(x,+1,mu,i,1)*U(x+mu,-1,nu,1,j)+U(x,+1,mu,i,2)*U(x+mu,-1,nu,2,j);
	  A3[i][j]=U(x,-1,mu,i,0)*U(x-mu,+1,nu,0,j)+U(x,-1,mu,i,1)*U(x-mu,+1,nu,1,j)+U(x,-1,mu,i,2)*U(x-mu,+1,nu,2,j);
	  A4[i][j]=U(x,-1,mu,i,0)*U(x-mu,-1,nu,0,j)+U(x,-1,mu,i,1)*U(x-mu,-1,nu,1,j)+U(x,-1,mu,i,2)*U(x-mu,-1,nu,2,j);
	}
      for(a=0; a<4; a++)
	for(i=0; i<3; i++)
	  for(j=0; j<3; j++) 
	    sink_n_shifted(x,a,i)=0.25*(A1[i][j]*sink_n(y1,a,j)-
					A2[i][j]*sink_n(y2,a,j)-
					A3[i][j]*sink_n(y3,a,j)+
					A4[i][j]*sink_n(y4,a,j));
      
    }	      
  }
}

int main(int argc, char** argv) {
  mdp.open_wormholes(argc,argv);
  define_base_matrices("FERMILAB");

  char filename_gauge[1024]="cold";
  if(argc>2) strcpy(filename_gauge,argv[1]);
  char filename_n[1024];
  char filename_m[1024];
  char filename_tmp[1024];
  char filename_output[1024];
  int nt=32,ns=4;
  int nstochastic=10;
  coefficients light_quark;
  light_quark["kappa_t"]=0.11;
  light_quark["kappa_s"]=0.11;
  light_quark["c_{sw}"]=0.1;
  float ap=1e-6;
  float rp=1e-6;
  default_fermi_inverter=MinimumResidueInverter<fermi_field,gauge_field>;
#ifdef SSE2
  default_fermi_action=FermiCloverActionSSE2::mul_Q;
#else
  default_fermi_action=FermiCloverActionFast::mul_Q;
#endif
  int box[]={nt,ns,ns,ns};
  mdp_lattice lattice(4,box);
  gauge_field U(lattice,3);
  fermi_field source_n(lattice,3);
  fermi_field source_m(lattice,3);
  fermi_field sink_n(lattice,3);
  fermi_field sink_m(lattice,3);  
  mdp_site x(lattice);
  mdp_site y(lattice);
  int i,j,k,n,m,a,b,t,ipath,ip;
  float p[4];
  mdp_complex theta;
  mdp_complex path[3][3];

  //
  // auxiliary variables
  //
  fermi_field sink_n_shifted(lattice,3);  

  mdp_lattice lattice_time(1,box);
  mdp_field<mdp_complex> pi_source(lattice_time,2*2*2);
  mdp_field<mdp_complex> pi_sink(lattice_time,21*4*4);
  mdp_site xp(lattice_time);
  mdp_site yp(lattice_time);
  mdp_complex q[nt][2*2*2][21][4][4];
  int q_size=nt*(2*2*2)*(21)*4*4*sizeof(mdp_complex);
  // 
  // set fields to zero
  //
  pi_source=0;
  pi_sink=0;
  for(t=0; t<nt; t++)
    for(ip=0; ip<2*2*2; ip++)
      for(ipath=0; ipath<21; ipath++)
	for(a=0; a<4; a++)
	  for(b=0; b<4; b++)
	    q[t][ip][ipath][a][b]=0;   
  
  //
  // read input gauge configuration
  //
  if(string(filename_gauge)==string("cold"))
    set_cold(U);
  else
    U.load(filename_gauge);
  compute_em_field(U);

  //
  // create stochastic propagators
  //
  for(n=0; n<nstochastic; n++) {
    generate_stochastic(source_n,U,light_quark,ap,rp);
    sprintf(filename_n,"%s_stochastic_source_%.3i", filename_gauge,n);
    source_n.save(filename_n);
    mul_Q(sink_n,source_n,U,light_quark);
    forallsites(x)
      for(a=0; a<4; a++)
	for(i=0; i<3; i++)
	  sink_n(x,a,i)=conj(sink_n(x,a,i));
    sprintf(filename_n,"%s_stochastic_sink_%.3i", filename_gauge,n);
    sink_n.save(filename_n);
  }


 
  for(n=0; n<nstochastic; n++) {
      
      // read stochastic prop n
    sprintf(filename_n,"%s_stochastic_source_%.3i", filename_gauge,n);
    source_n.load(filename_n);
    sprintf(filename_n,"%s_stochastic_sink_%.3i", filename_gauge,n);
    sink_n.load(filename_n);

    // ipath =0    => 1
    // ipath =1+mu => D_mu
    // ipath =5+4*mu+nu => D_mu D_nu
    for(ipath=0; ipath<21; ipath++) {	 	
      build_shifted(sink_n_shifted,U,sink_n,ipath);
      sprintf(filename_tmp,"tmp_%.3i",ipath);
      sink_n_shifted.save(filename_tmp);
    }   

    for(m=0; m<nstochastic; m++) {
    
    // read stochastic prop m
      sprintf(filename_m,"%s_stochastic_source_%.3i", filename_gauge,m);
      source_m.load(filename_m);
      sprintf(filename_m,"%s_stochastic_sink_%.3i", filename_gauge,m);
      sink_m.load(filename_m);
      
      // make pion sources fft
      
      for(ip=0; ip<2*2*2; ip++) {
	p[1]=(ip>>2 & 0xf);
	p[2]=(ip>>1 & 0xf);
	p[3]=(ip>>0 & 0xf);
	forallsites(x) {
	  theta=exp(2.0*Pi*I*(p[1]*x(1)/box[1]+
			      p[2]*x(2)/box[2]+
			      p[3]*x(3)/box[3]))/lattice.global_volume();
	  xp.set(x(0));
	  for(a=0; a<4; a++)
	    for(i=0; i<3; i++)
	      pi_source(xp,ip)+=source_m(x,a,i)*conj(source_n(x,a,i))*theta;
	}
      }
      //
      // make pion sink for each x and each path
      //
      /*
      for(ipath=0; ipath<21; ipath++) {	 

	sprintf(filename_tmp,"tmp_%.3i",ipath);
	sink_n_shifted.load(filename_tmp);
	
	forallsites(x) {	 
	  xp.set(x(0));
	  for(a=0; a<4; a++)
	    for(b=0; b<4; b++) {
		pi_sink(xp,ipath*16+a*4+b)+=
		  (conj(sink_m(x,a,0))*sink_n_shifted(x,b,0)+
		   conj(sink_m(x,a,1))*sink_n_shifted(x,b,1)+
		   conj(sink_m(x,a,2))*sink_n_shifted(x,b,2))/lattice.global_volume();
	    }
	}
      }
      */
      forallsites(xp) {
	forallsites(yp) {
	  t=(nt+xp(0)-yp(0)) % nt;	     
	  for(ip=0; ip<2*2*2; ip++) 
	    for(ipath=0; ipath<21; ipath++) 
	      for(a=0; a<4; a++)
		for(b=0; b<4; b++) {
		  q[t][ip][ipath][a][b]+=
		    pi_source(xp,ip)*
		    pi_sink(yp,ipath*16+a*4+b)/pow(nstochastic,2);
		}	 
	}
	
      }
    }
  }

  for(t=0; t<nt; t++)
    cout << t 
	 << "\t" <<  real(q[t][0][0][0][0]+q[t][0][0][1][1]+q[t][0][0][2][2]+q[t][0][0][3][3])
	 << "\t" <<  real(q[t][0][1][0][0]+q[t][0][1][1][1]+q[t][0][1][2][2]+q[t][0][1][3][3])
	 << "\t" <<  real(q[t][0][2][0][0]+q[t][0][2][1][1]+q[t][0][2][2][2]+q[t][0][2][3][3])
	 << "\t" <<  real(q[t][0][3][0][0]+q[t][0][3][1][1]+q[t][0][3][2][2]+q[t][0][3][3][3]) << endl;
  
  sprintf(filename_output,"%s_qarray", filename_gauge);
  FILE *fp=fopen(filename_output,"wb");
  fwrite(q,q_size,1,fp);
  fclose(fp);
  
  
  mdp.close_wormholes();
}
  
